from Final_assignment.PartitionDataset import *
from Final_assignment.generate_Kline_diagram import genenate_Kline_diagram
from Final_assignment.tools import read_file_list
from Final_assignment.train import *


# 程序测试模式
def test_mode():
    stock_list_test = read_file_list("files/", "stock_list_test.txt")
    genenate_Kline_diagram("weekly", stock_list_test, "test")  # 下载数据并绘制K线图
    print("genenate Kline diagram successfully")
    partition_dataset("weekly", "train_dataset_train", "test_dataset_test", stock_list_test)  # 划分数据集，训练集和测试集
    print("partition dataset successfully (test mode)")
    model = only_train("train_dataset_train")  # 训练
    print("train successfully (test mode)")
    only_test("model/model.pkl", "test_dataset_test", model)  # 测试
    print("test successfully (test mode)")


if __name__ == "__main__":
    # 划分数据集，周期，训练集文件夹，测试集文件夹
    mode = input(
        "input the mode of program running ('test'/other string),\n\
        'test' will run program quickly,but not accuracy,other input is run program normally:")
    train_folder = "train_dataset2"  # 训练集文件夹
    test_folder = "test_dataset2"  # 测试集文件夹

    if mode == "test":  # mode = "test"
        test_mode()  # 执行测试模式
    elif mode == '1':
        train_folder = "train_dataset1"  # 训练集文件夹
        test_folder = "test_dataset1"  # 测试集文件夹
        stock_list = read_file_list("files/", "stock_list1.txt")  # 读取股票列表,上证50，50只
    elif mode == '2':
        train_folder = "train_dataset2"  # 训练集文件夹
        test_folder = "test_dataset2"  # 测试集文件夹
        stock_list = read_file_list("files/", "stock_list2.txt")  # 读取股票列表,银行股，29只

    if mode != "test":
        genenate_Kline_diagram("weekly", stock_list, mode)  # 下载数据并绘制K线图
        print("genenate Kline diagram successfully")

        partition = input("please input 'yes' to partition dataset again,and other input to ignore this:")
        if partition == "yes":
            partition_dataset("weekly", train_folder, test_folder, stock_list)  # 划分数据集，训练集和测试集
            print("partition dataset successfully")
        model = only_train(train_folder)  # 训练
        print("train successfully")
        only_test("model/model.pkl", test_folder, model)  # 测试
        print("test successfully")
